# 向量存储类

import os
import chromadb
from datetime import datetime
from .model_loader import model_loader
from .models import Document, Category
from config import (
    CHROMA_PERSIST_DIR,
)


class VectorStore:
    """向量存储类，用于文档向量的存储和相似度搜索"""

    def __init__(self, db):
        """初始化向量存储"""
        self.db = db
        # 确保ChromaDB持久化目录存在
        os.makedirs(CHROMA_PERSIST_DIR, exist_ok=True)
        # 初始化ChromaDB客户端，使用持久化存储
        self.client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR)
        # 创建或获取集合
        self.collection = self.client.get_or_create_collection("documents")

    def get_text_embedding(self, text: str) -> list[float]:
        """获取文本的嵌入向量"""
        # 统一通过model_loader获取嵌入向量，model_loader会根据API_TYPE选择合适的方式
        return model_loader.generate_embedding(text)

    def store_document(self, content, category_ids=None, title=None):
        """存储文档及其向量嵌入，支持多个分类ID"""
        session = self.db.connect()
        try:
            # 1. 创建文档并添加到会话
            document = Document(title=title, content=content)
            session.add(document)
            session.commit()
            # 刷新会话以确保获取完整的文档对象
            session.refresh(document)

            # 2. 如果提供了分类ID，将文档添加到分类
            if category_ids:
                # 获取所有有效的分类
                categories = session.query(Category).filter(Category.id.in_(category_ids)).all()
                if categories:
                    document.categories.extend(categories)
                    session.commit()

            # 3. 存储向量到ChromaDB
            # 将标题和内容合并后进行嵌入，使标题也成为检索依据
            text_to_embed = content
            if title:
                text_to_embed = f"{title}\n{content}"

            embedding = self.get_text_embedding(text_to_embed)
            # 为向量添加元数据，将分类ID列表转换为逗号分隔的字符串
            metadata = {
                "document_id": document.id,
                "category_ids": ",".join(map(str, category_ids)) if category_ids else "",
                "title": title if title else "",
                "content": content,
            }
            self.collection.add(ids=[str(document.id)], embeddings=[embedding], metadatas=[metadata])

            return document.id
        finally:
            # 确保在访问完所有需要的属性后再关闭会话
            session.close()

    def search_similar_documents(self, query_text, top_k=5, category_id=None):
        """搜索相似文档"""
        # 获取查询文本的向量嵌入
        query_embedding = self.get_text_embedding(query_text)

        # 构建过滤条件
        where_clause = None
        if category_id:
            # 我们需要先查询数据库获取所有属于该分类的文档ID
            session = self.db.connect()
            try:
                # 查找属于该分类的所有文档
                category = session.query(Category).filter(Category.id == category_id).first()

                if category:
                    # 获取属于该分类的所有文档ID，保持整数类型
                    document_ids = [doc.id for doc in category.documents]

                    if document_ids:
                        # 使用$in操作符来过滤文档ID，确保类型匹配
                        where_clause = {"document_id": {"$in": document_ids}}
            finally:
                session.close()

        # 执行相似度搜索
        results = self.collection.query(query_embeddings=[query_embedding], n_results=top_k, where=where_clause)

        # 处理搜索结果
        search_results = []
        session = self.db.connect()
        try:
            for i in range(len(results["ids"][0])):
                doc_id = int(results["ids"][0][i])
                metadata = results["metadatas"][0][i]
                similarity = results["distances"][0][i]

                # 从数据库获取文档的分类信息
                document = session.query(Document).filter(Document.id == doc_id).first()
                category_ids = [cat.id for cat in document.categories] if document else []

                # 处理相似度计算，添加异常值处理和归一化
                # 对于Doubao-embedding，距离可能非常大（负值），需要特殊处理
                try:
                    # 检查是否为异常值
                    if similarity < 0:
                        # 对于负值，可以直接取绝对值并进行归一化
                        # 假设合理距离范围在0-100之间
                        normalized_score = max(0, min(1, 1 - abs(similarity) / 100))
                    elif similarity > 100:
                        # 对于特别大的正值，也进行归一化处理
                        normalized_score = max(0, 1 - similarity / 1000)
                    else:
                        # 正常范围内的值，按原逻辑处理
                        normalized_score = 1 - similarity
                except:
                    # 异常情况下默认为0
                    normalized_score = 0
                
                # 确保分数在0-1之间
                normalized_score = max(0, min(1, normalized_score))
                
                # 转换为与原始接口兼容的格式，添加分类ID信息
                search_results.append(
                    (
                        doc_id,
                        metadata["content"],
                        metadata["title"] if metadata["title"] else None,
                        normalized_score,
                        category_ids,
                    )
                )
        finally:
            session.close()

        return search_results
